from contextlib import closing
from typing import Optional
from db_exporter.exception import DBExportException
from .base import BasePlugin


class DbApi(BasePlugin):
    _test_connection_sql = "select 1"
    conn_name_attr = "conn_id"

    def __init__(self, *args, schema: Optional[str] = None, **kwargs):
        super().__init__()
        if not self.conn_name_attr:
            raise DBExportException("conn_name_attr is not defined")
        elif len(args) == 1:
            setattr(self, self.conn_name_attr, args[0])
        elif self.conn_name_attr not in kwargs:
            setattr(self, self.conn_name_attr, self.default_conn_name)
        else:
            setattr(self, self.conn_name_attr, kwargs[self.conn_name_attr])
        self.__schema = schema

    def get_conn(self):
        """Returns a connection object"""
        db = self.get_connection(getattr(self, self.conn_name_attr))
        return self.connector.connect(host=db.host, port=db.port, username=db.login, schema=db.schema)

    def get_uri(self) -> str:
        """
        Extract the URI from the connection.

        :return: the extracted uri.
        """
        conn = self.get_connection(getattr(self, self.conn_name_attr))
        conn.schema = self.__schema or conn.schema
        return conn.get_uri()

    def get_cursor(self):
        """Returns a cursor"""
        return self.get_conn().cursor()

    def _get_cursor_by_type(self, conn, cursor=None):
        return conn.cursor()

    def get_records(self, sql, parameters=None, cursor=None):
        """
        Executes the sql and returns a set of records.

        :param sql: the sql statement to be executed (str) or a list of
            sql statements to execute
        :param parameters: The parameters to render the SQL query with.
        """
        with closing(self.get_conn()) as conn:
            with closing(self._get_cursor_by_type(conn, cursor)) as cur:
                if parameters is not None:
                    cur.execute(sql, parameters)
                else:
                    cur.execute(sql)
                return cur.fetchall()

    def get_records_md5(self, sql, parameters=None, cursor=None):
        """Returns the md5 list of records """
        raise NotImplementedError()

    def get_cols(self, table_name: str):
        with closing(self.get_conn()) as conn:
            with closing(conn.cursor()) as cur:
                sql = "select * from (%s) where 1=2" % table_name
                cur.execute(sql)
                col = cur.description

                cols = []
                for i in range(len(col)):
                    cols.append(col[i][0])

                return cols

    def get_cols_desc(self, sql: str):
        """

        :param sql:
        :return:
        list of col description. (col_name, col_type, None, total_len, None, scale, null_ok)

            col_type, 参见 pymysql.constants.FIELD_TYPE.py 246:NEWDECIMAL, 253:varchar

        """
        with closing(self.get_conn()) as conn:
            with closing(conn.cursor()) as cur:
                sql = "select * from (%s) t where 1=2" % sql
                cur.execute(sql)
                cols_desc = cur.description

                return cols_desc

    def get_cols_by_sql(self, sql: str):
        with closing(self.get_conn()) as conn:
            with closing(conn.cursor()) as cur:
                sql = "select * from (%s) t where 1=2" % sql
                cur.execute(sql)
                col = cur.description

                cols = []
                for i in range(len(col)):
                    cols.append(col[i][0])

                return cols

    def get_count(self, sql: str):
        cnt_sql = "select count(1) from (%s) l1" % sql
        return self.get_first(cnt_sql)[0]

    def get_first(self, sql, parameters=None):
        """
        Executes the sql and returns the first resulting row.

        :param sql: the sql statement to be executed (str) or a list of
            sql statements to execute
        :param parameters: The parameters to render the SQL query with.
        """
        with closing(self.get_conn()) as conn:
            with closing(conn.cursor()) as cur:
                try:
                    if parameters is not None:
                        cur.execute(sql, parameters)
                    else:
                        cur.execute(sql)
                    return cur.fetchone()
                except Exception as e:
                    print("conn_id:%s, sql:%s, parameters:%s" % (self.conn_id, sql, parameters))
                    raise e

    def test_connection(self):
        """Tests the connection using db-specific query"""
        status, message = False, ''
        try:
            if self.get_first(self._test_connection_sql):
                status = True
                message = 'Connection successfully tested'
        except Exception as e:
            status = False
            message = str(e)

        return status, message
